import os
import shutil
import numpy as np

class MiniDatasets:
    def __init__(self):
        super(MiniDatasets, self).__init__()
        train_img_dir = 'data/afhq/train'
        val_img_dir = 'data/afhq/val'
        mini_datasets_ratio = 0.1
        self.generator_mini_datasets(train_img_dir, 'data/mini_afhq/train', mini_datasets_ratio)
        self.generator_mini_datasets(val_img_dir, 'data/mini_afhq/val', mini_datasets_ratio)

    def generator_mini_datasets(self, img_dir, save_dir, datasets_ratio=0.1):
        domains = [d.name for d in os.scandir(img_dir)]
        print('domains:', domains)

        for domain in domains:
            sub_img_dir = os.path.join(img_dir, domain)
            sub_save_dir = os.path.join(save_dir, domain)
            if os.path.exists(sub_save_dir):
                shutil.rmtree(sub_save_dir)
            os.makedirs(sub_save_dir, exist_ok=True)
            imgs = [d.name for d in os.scandir(sub_img_dir)]
            np.random.shuffle(imgs)
            count = int(len(imgs) * datasets_ratio)
            new_imgs = imgs[:count]
            print('domain:', domain, 'imgs:', len(imgs), 'count:', count, 'new_imgs:', len(new_imgs))
            for img in new_imgs:
                shutil.copyfile(os.path.join(sub_img_dir, img), os.path.join(sub_save_dir, img))


if __name__ == '__main__':
    mini_datasets = MiniDatasets()

